from PIL import Image


def nonlinear_sharpening_filter_rgb(image_path, output_path):
    """
    对输入的RGB图像进行非线性锐化滤波处理。

    参数:
    image_path (str): 输入图像的文件路径。
    output_path (str): 输出滤波后图像的文件路径。
    """
    # 打开原始图像
    image = Image.open(image_path)
    width, height = image.size
    # 创建新的空白图像，尺寸与原图像相同，模式也相同，用于存储锐化后的图像
    new_image = Image.new(image.mode, (width, height))

    # 定义局部区域大小（这里以3x3为例，可根据实际调整）
    neighborhood_size = 3
    radius = neighborhood_size // 2

    # 遍历图像像素点进行非线性锐化处理（避开边缘部分，可后续优化边缘处理）
    for y in range(radius, height - radius):
        for x in range(radius, width - radius):
            # 用于累计各通道的局部像素值总和
            sum_red = 0
            sum_green = 0
            sum_blue = 0
            # 用于累计各通道的局部像素值平方总和
            sum_square_red = 0
            sum_square_green = 0
            sum_square_blue = 0
            # 遍历局部区域内的像素点
            for ky in range(-radius, radius + 1):
                for kx in range(-radius, radius + 1):
                    pixel = image.getpixel((x + kx, y + ky))
                    red, green, blue = pixel
                    sum_red += red
                    sum_green += green
                    sum_blue += blue
                    sum_square_red += red ** 2
                    sum_square_green += green ** 2
                    sum_square_blue += blue ** 2

            # 计算局部区域内各通道的像素均值
            num_pixels = neighborhood_size ** 2
            mean_red = sum_red // num_pixels
            mean_green = sum_green // num_pixels
            mean_blue = sum_blue // num_pixels

            # 计算局部区域内各通道的方差
            variance_red = (sum_square_red // num_pixels) - (mean_red ** 2)
            variance_green = (sum_square_green // num_pixels) - (mean_green ** 2)
            variance_blue = (sum_square_blue // num_pixels) - (mean_blue ** 2)

            # 获取原图像当前像素点的像素值
            original_pixel = image.getpixel((x, y))
            red_original, green_original, blue_original = original_pixel

            # 根据方差计算锐化系数（经实验将方差除以2000再加一效果较好，加一是因为锐化系数要大于1）
            k_red = 1 + variance_red / 2000
            k_green = 1 + variance_green / 2000
            k_blue = 1 + variance_blue / 2000

            # 计算锐化后的像素值（限制在0-255范围内）
            red_sharpened = max(0, min(255, int(red_original * k_red)))
            green_sharpened = max(0, min(255, int(green_original * k_green)))
            blue_sharpened = max(0, min(255, int(blue_original * k_blue)))

            # 将锐化后的像素值设置到新图像对应坐标位置
            new_image.putpixel((x, y), (red_sharpened, green_sharpened, blue_sharpened))

    # 保存锐化后的图像
    new_image.save(output_path)
    return new_image

image_path = "D:/cangku/computer-image-project-design/test.jpg"
output_path = "D:/cangku/computer-image-project-design/test1.jpg"
nonlinear_sharpening_filter_rgb(image_path, output_path)